import os
import json
import time
import gradio as gr
import pandas as pd
import xml.etree.ElementTree as ET
from datetime import datetime
import dateparser
import dashscope

# 初始化API密钥
dashscope.api_key = os.getenv('DASHSCOPE_API_KEY')


def normalize_time(time_str):
    """标准化时间格式（英语）"""
    if not time_str or time_str.lower() in ['n/a', 'unknown']:
        return 'N/A'

    try:
        settings = {'DATE_ORDER': 'MDY'}
        dt = dateparser.parse(time_str, settings=settings)
        return dt.isoformat() if dt else time_str
    except:
        return time_str


def extract_events(story):
    """使用大模型提取事件四元组"""
    try:
        messages = [
            {
                'role': 'system',
                'content': """Extract story events as JSON array with:
- head: Subject (person/entity)
- relation: Action (past tense)
- tail: Object/target
- time: Normalized timestamp (or N/A)"""
            },
            {
                'role': 'user',
                'content': f"""Story: {story}

Respond ONLY with JSON array. Example:
[
  {{"head": "Cinderella", "relation": "was forced to work for", 
    "tail": "stepmother", "time": "after father's death"}}
]"""
            }
        ]

        response = dashscope.Generation.call(
            model='qwen-max',
            messages=messages,
            result_format='message'
        )

        events = json.loads(response.output.choices[0].message.content)
        for event in events:
            event['time'] = normalize_time(event.get('time', 'N/A'))
        return events
    except Exception as e:
        print(f"Extraction Error: {str(e)}")
        return []


def generate_premise(events):
    """根据事件生成故事前提"""
    try:
        if not events:
            return "⚠️ No events available"

        events_text = "\n".join(
            [f"{e['head']} {e['relation']} {e['tail']} ({e['time']})"
             for e in events]
        )

        messages = [
            {
                'role': 'system',
                'content': """Generate a 2-sentence story premise connecting these events. 
                Include suspense and logical flow."""
            },
            {
                'role': 'user',
                'content': f"Events:\n{events_text}\n\nStory Premise:"
            }
        ]

        response = dashscope.Generation.call(
            model='qwen-max',
            messages=messages
        )
        return response.output.choices[0].message.content
    except Exception as e:
        return f"❌ Generation failed: {str(e)}"


def save_to_xml(events, filename):
    """保存事件到XML文件"""
    root = ET.Element('story_events')

    # 元数据
    meta = ET.SubElement(root, 'metadata')
    ET.SubElement(meta, 'created').text = datetime.now().isoformat()
    ET.SubElement(meta, 'language').text = 'en'
    ET.SubElement(meta, 'source').text = filename

    # 事件数据
    events_elem = ET.SubElement(root, 'events')
    for event in events:
        event_elem = ET.SubElement(events_elem, 'event')
        ET.SubElement(event_elem, 'head').text = event.get('head', '')
        ET.SubElement(event_elem, 'relation').text = event.get('relation', '')
        ET.SubElement(event_elem, 'tail').text = event.get('tail', '')
        ET.SubElement(event_elem, 'time').text = event.get('time', 'N/A')

    # XML格式化
    def indent(elem, level=0):
        spacer = "  "
        if len(elem):
            if not elem.text or not elem.text.strip():
                elem.text = "\n" + spacer * (level + 1)
            for child in elem:
                indent(child, level + 1)
            if not elem.tail or not elem.tail.strip():
                elem.tail = "\n" + spacer * level
        else:
            if level and (not elem.tail or not elem.tail.strip()):
                elem.tail = "\n" + spacer * level

    indent(root)

    os.makedirs("stories", exist_ok=True)
    full_path = os.path.join("stories", filename)
    ET.ElementTree(root).write(full_path, encoding='utf-8', xml_declaration=True)
    return full_path


def process_story(story_text):
    """处理故事分析流程"""
    start_time = time.time()

    try:
        # 事件提取
        events = extract_events(story_text)
        if not events:
            return [], "⚠️ No events extracted", "", []

        # 保存文件
        timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
        filename = f"story_{timestamp}.xml"
        saved_path = save_to_xml(events, filename)

        # 准备输出
        process_time = f"⏱️ {time.time() - start_time:.2f}s"
        return events, process_time, saved_path, events
    except Exception as e:
        return [], f"❌ Error: {str(e)}", "", []


# 界面布局
css = """
.dataframe {max-height: 300px !important;}
.markdown-textarea {min-height: 100px !important;}
footer {visibility: hidden}
"""

with gr.Blocks(theme=gr.themes.Soft(), css=css, title="Story Architect") as app:
    gr.Markdown("# 📖 Story Architect - Interactive Narrative Builder")

    with gr.Row():
        with gr.Column(scale=3):
            input_box = gr.Textbox(label="Input Story", lines=8,
                                   placeholder="Paste your story here...")
            analyze_btn = gr.Button("Analyze Story", variant="primary")

        with gr.Column(scale=4):
            with gr.Tab("Analysis"):
                original_table = gr.DataFrame(
                    label="Extracted Events",
                    headers=["head", "relation", "tail", "time"],
                    datatype=["str", "str", "str", "str"],
                    interactive=False
                )
                file_output = gr.Textbox(label="Saved Path", visible=False)

            with gr.Tab("Edit & Generate"):
                editable_table = gr.DataFrame(
                    label="Edit Events (modify values below)",
                    headers=["head", "relation", "tail", "time"],
                    datatype=["str", "str", "str", "str"],
                    interactive=True
                )
                gen_btn = gr.Button("Generate New Premise", variant="secondary")
                premise_output = gr.Textbox(label="Generated Story Premise", lines=3)

            process_info = gr.Textbox(label="Processing Status", interactive=False)

    # 示例和交互逻辑
    gr.Examples(
        examples=[[
            "In a futuristic city, young hacker Zane discovered a secret AI program. "
            "He uploaded the AI to his neural interface, gaining extraordinary hacking abilities. "
            "The corrupt city mayor launched a city-wide manhunt to capture Zane."
        ]],
        inputs=input_box,
        outputs=[original_table, process_info, file_output, editable_table],
        fn=process_story
    )

    analyze_btn.click(
        fn=process_story,
        inputs=input_box,
        outputs=[original_table, process_info, file_output, editable_table]
    )

    gen_btn.click(
        fn=generate_premise,
        inputs=editable_table,
        outputs=premise_output
    )

if __name__ == "__main__":
    app.launch(server_name="127.0.0.1", server_port=7860)